#pragma once

#include <NvInfer.h>
#include <NvInferRuntimeCommon.h>

#include <memory>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

#include "nn_common.hpp"

using namespace nvinfer1;

enum class CreateMethod {
  kCACHE = 0,
  kWEIGHTS = 1,
  kONNX = 2,
  kFILE = 3,
};

class TRTYoloBuilder {
 private:
  std::string file_path_;
  std::string engine_path_;

  TRT::TRTLogger logger_;

  UniquePtr<ICudaEngine> engine_;
  UniquePtr<IExecutionContext> context_;
  UniquePtr<INetworkDefinition> network_;

  bool CreateEngineWithOnnx(IBuilder* builder, IBuilderConfig* config);
  bool CreateEngineWithWeights(IBuilder* builder, IBuilderConfig* config,
                               DataType datatype);
  bool CreateEngineWithCache();
  bool CreateEngineWithFile();

 public:
  TRTYoloBuilder();
  TRTYoloBuilder(const std::string& file_path, const std::string& engine_path);
  ~TRTYoloBuilder();

  void SetFilePath(const std::string& file_path);
  void SetEnginePath(const std::string& engine_path);
  bool CreateEngine(CreateMethod method);
  bool SaveEngine();

  UniquePtr<ICudaEngine> GetEngine();
};